#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
import src.MRTools as mrtools
from inspect import stack

class Lockin_Detection():
    def __init__(self,fundamental_f=100000, n_harmonics=1):
        """
        Inputs:
            -fundamental_f: in Hz
            -n_harmonics: integer, 1-> 2*f0, 2-> 3*f0
        """
        self.f0=fundamental_f
        self.harmonics=np.arange(1,n_harmonics+2,1)      
        self.signals_X={}
        self.signals_Y={}  
        self.signals_R={}
        self.signals_phi={}

        
    def detect_XY(self,time_ax,signal):
        """
        Input:
            -time_ax: the corresponding time axis, np.array([n])
            -signal: np.array([n])
        """
        for harm in self.harmonics:
            f_det=harm*self.f0
            phi0 = np.sum(signal*np.sin(2*np.pi*f_det*time_ax))
            phi90 = np.sum(signal*np.sin(2*np.pi*f_det*time_ax+np.pi/2))
            self.signals_X[harm]=phi0
            self.signals_Y[harm]=phi90

    def calculate_Rphi(self):
        for harm in self.harmonics:
            self.signals_R[harm]=np.sqrt(self.signals_X[harm]**2+self.signals_Y[harm]**2)
            # arctan2(y,x)
            self.signals_phi[harm]=np.arctan2(self.signals_Y[harm],self.signals_X[harm])

if __name__ == "__main__":
    def func(t,A,tau):
        return A*(np.exp(-t/tau))
    
    step=0.000001
    t=np.arange(0,1,step)
    s=np.sin(2*np.pi*100000*t+np.pi/3)*func(t,1,1)
    
    #lockin=Lockin_Detection()
    #lockin.detect_XY(t,s)
    #lockin.calculate_Rphi()
    
    fig = plt.figure()
    ax=fig.add_subplot(1,1,1)
    ax.plot(np.fft.fftfreq(s.size, step),np.abs(np.fft.fft(s)),'b')
    plt.show()